import networkx as nx
import util
import types
from itertools import izip

from core.std import Component


class Edge(object):
    def __init__(self, out_node, out_attr, in_node, in_attr):
        self.out_node = out_node
        self.out_attr = out_attr
        self.in_node = in_node
        self.in_attr = in_attr
        
    def out_data(self):
        return getattr(self.out_node, self.out_attr)
    
    def in_data(self):
        return getattr(self.in_node, self.in_attr)


class Pipeline(object):
    
    def __init__(self, *args, **kw):
        self.graph = nx.MultiDiGraph()
    
    def add_node(self, node):
        self.graph.add_node(node, dirty=True)
    
    def add_nodes(self, *nodes):
        map(self.add_node, nodes)
        
    def nodes(self, data=None):
        return self.graph.nodes(data)
    
    def edges(self, attrs=None):
        for out_node, in_node, data in self.graph.edges(data=True):
            if attrs is None:
                yield (out_node, in_node)
                continue
            out_attr = out_node.get_attr(data['out_attr'])
            in_attr = in_node.get_attr(data['in_attr'])
            yield (out_node, out_attr, in_node, in_attr)
        
    def connect(self, out_node, out_attr, in_node, in_attr):
        self.graph.add_edge(out_node, in_node, in_attr=in_attr, out_attr=out_attr)
    
    def topological_sort_iter(self):
        for node in nx.topological_sort(self.graph):
            yield node
    
    def update(self):
        for node in self.graph.nodes():
            node.prepare_update()
        root_node = list(self.topological_sort_iter())[0]
        if root_node is None: return
        self.update_node(root_node)
        self.dirty_all_nodes()
        
    def dirty_all_nodes(self):
        for node in self.nodes():
            self.dirty_node(node)
            
    def dirty_node(self, node):
        self.graph.node[node]['dirty'] = True
        
    def dirty_dependency_nodes(self, root):
        for node in nx.dfs_successor(self.graph, root).keys()[1:]:
            self.graph.node[node]['dirty'] = True
        
    def update_node(self, node):
        if not self.graph.node[node]['dirty']: 
            return
        self.pull_data(node)
    
        node.update()
        self.graph.node[node]['dirty'] = False
        for attr in node.output_attrs():
            data = getattr(node, attr.name)
            attr_edges = self.out_edges(node, attr)
            if isinstance(data, types.GeneratorType):
                for item in getattr(node, attr.name):
                    self.dirty_dependency_nodes(node)
                    for edge in self.out_edges(node, attr):
                        self.push_data(edge, item)
                    for next_node in self.next(node):
                        self.update_node(next_node)
            else:
                for next_node in self.next(node):
                    for edge in attr_edges:
                        self.push_data(edge, data)
                    self.update_node(next_node)
        
    def graphviz_layout(self):
        return nx.graphviz_layout(self.graph)
    
    def depth_node_iter(self, root=None, depth=0):
        for child in self.next(root):
            yield (depth, child)
            for depth, sub_child in self.depth_node_iter(child, depth):
                yield (depth+1, sub_child)
        
    def tree_layout(self):
        root = list(self.topological_sort_iter())[0]
        tree = []
        tree.append([root])
        pos_dict = {}
        
        for depth, child in self.depth_node_iter(root):
            if not len(tree) > depth+1:
                tree.append([])
            tree[depth+1].append(child)
        
        x,y = 0,0
        xspace, yspace = 150, -40
        for depth, nodes in enumerate(tree):
            for node in nodes:
                pos_dict[node] = (depth*xspace, len(nodes)*yspace,)
                
        return pos_dict
        
    def pull_data(self, node):
        for dirty_node in self.prev(node, dirty=True):
            self.update_node(dirty_node)
    
    def push_data(self, edge, data):
        setattr(edge.in_node, edge.in_attr, data)
    
    def next(self, node, dirty=None):
        if dirty is None:
            return self.graph.successors(node)
        return [n for n in self.graph.successors(node) if self.graph.node[n]['dirty'] == dirty]
    
    def prev(self, node, dirty=None):
        if dirty is None:
            return self.graph.predecessors(node)
        return [n for n in self.graph.predecessors(node) if self.graph.node[n]['dirty'] == dirty]
                
    def out_edges(self, node, out_attr=None):
        for out_node, in_node, data in self.graph.out_edges_iter(node, data=True):
            if out_attr is None or data['out_attr'] == out_attr.name:
                yield Edge(out_node, data['out_attr'], in_node, data['in_attr'])
    
